from typing import List, Dict, Any, Optional, Tuple
# from transformers import pipeline
from loguru import logger
import torch
import numpy as np
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import IsolationForest

class AIDecisionMaker:
    def __init__(self):
        # 初始化各种AI模型
        # self.classifier = pipeline("zero-shot-classification", 
        #                          model="joeddav/xlm-roberta-large-xnli")
        # self.text_generator = pipeline("text-generation",
        #                              model="THUDM/chatglm2-6b",
        #                              trust_remote_code=True)
        
        # 异常检测模型
        self.anomaly_detector = IsolationForest(contamination=0.1, random_state=42)
        self.scaler = StandardScaler()
        
        logger.info("AI决策模块初始化完成")

    def classify_task(self, text: str, labels: List[str]) -> Dict[str, float]:
        """任务分类"""
        try:
            # result = self.classifier(text, labels)
            # logger.info(f"任务分类完成: {result['labels'][0]}")
            # return dict(zip(result["labels"], result["scores"]))
            return {label: 1.0 / len(labels) for label in labels}  # 临时返回均匀分布
        except Exception as e:
            logger.error(f"任务分类失败: {str(e)}")
            raise

    def generate_action_plan(self, task_description: str) -> str:
        """生成行动计划"""
        try:
            # prompt = f"基于以下任务描述生成详细的自动化操作步骤：\n{task_description}\n\n步骤："
            # result = self.text_generator(prompt, max_length=500, num_return_sequences=1)
            # plan = result[0]["generated_text"]
            plan = "1. 分析任务\n2. 执行操作\n3. 验证结果"  # 临时返回固定计划
            logger.info("行动计划生成完成")
            return plan
        except Exception as e:
            logger.error(f"生成行动计划失败: {str(e)}")
            raise

    def detect_anomaly(self, data: np.ndarray) -> np.ndarray:
        """异常检测"""
        try:
            # 数据标准化
            scaled_data = self.scaler.fit_transform(data)
            # 检测异常
            predictions = self.anomaly_detector.fit_predict(scaled_data)
            anomalies = predictions == -1
            logger.info(f"异常检测完成，发现 {sum(anomalies)} 个异常")
            return anomalies
        except Exception as e:
            logger.error(f"异常检测失败: {str(e)}")
            raise

    def optimize_sequence(self, actions: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """优化操作序列"""
        try:
            # 简单的基于规则的优化
            # 1. 合并连续的相似操作
            # 2. 删除冗余操作
            optimized = []
            prev_action = None
            
            for action in actions:
                if prev_action is None:
                    optimized.append(action)
                    prev_action = action
                    continue
                    
                if self._can_merge_actions(prev_action, action):
                    prev_action = self._merge_actions(prev_action, action)
                    optimized[-1] = prev_action
                else:
                    optimized.append(action)
                    prev_action = action
            
            logger.info(f"操作序列优化完成，从 {len(actions)} 步优化到 {len(optimized)} 步")
            return optimized
        except Exception as e:
            logger.error(f"操作序列优化失败: {str(e)}")
            raise

    def _can_merge_actions(self, action1: Dict[str, Any], action2: Dict[str, Any]) -> bool:
        """判断两个操作是否可以合并"""
        # 示例规则，可以根据实际需求扩展
        if action1["type"] == action2["type"]:
            if action1["type"] == "click" and self._positions_are_close(
                action1["position"], action2["position"]
            ):
                return True
            if action1["type"] == "type" and action1["target"] == action2["target"]:
                return True
        return False

    def _merge_actions(self, action1: Dict[str, Any], action2: Dict[str, Any]) -> Dict[str, Any]:
        """合并两个操作"""
        if action1["type"] == "click":
            # 使用第二个点击的位置
            return action2
        if action1["type"] == "type":
            # 合并输入文本
            return {
                "type": "type",
                "target": action1["target"],
                "text": action1["text"] + action2["text"]
            }
        return action2

    def _positions_are_close(self, pos1: Tuple[int, int], pos2: Tuple[int, int], threshold: int = 10) -> bool:
        """判断两个位置是否接近"""
        return abs(pos1[0] - pos2[0]) <= threshold and abs(pos1[1] - pos2[1]) <= threshold 